#!/usr/bin/python
import os
import shutil
import sys
import getopt
import time
import locale
import re
import datetime

import MySQLdb

from OVIMTools import _OVIMLog

def _get_filelist_byext(file_extlist, rootpath):                                 
    fileList = [os.path.normcase(f)
                for f in os.listdir(rootpath)]           
    fileList = [os.path.join(rootpath, f) 
               for f in fileList
                if os.path.splitext(f)[1] in file_extlist]
    return fileList

class VMWareHelper:
    _vmxpath = ''
    _vmxshellpath = ''
    _vmrootpath = ''
    _vmwarecmdpath = "/usr/bin/vmware-cmd"
    _vmwarerunpath = "/usr/bin/vmrun"
    _log = None
    _wait_timeout = 300

    def get_status(self):
        self._log.log("VMWareHelper:: try to get status", 'DEBUG')
        onoffStream = os.popen("%s \"%s\" getstate| cut -d \" \" -f 3" \
                                % (self._vmwarecmdpath, self._vmxpath))
        status = onoffStream.read().strip()
        self._log.log("VMWareHelper:: getStatus will return %s" % status, 'DEBUG')
        return status
    
    def take_snapshot(self):
        self._log.log("VMWareHelper:: try to take snapshot", 'DEBUG')
        self._log.log("Taking a snapshot")
        self._log.log("%s snapshot \"%s\"" \
                                % (self._vmwarerunpath, self._vmxpath), 'DEBUG')
        onoffStream = os.popen("%s snapshot \"%s\"" \
                                % (self._vmwarerunpath, self._vmxpath))
        status = onoffStream.read().strip()
        
        if status != '':
            self._log.log("VMWareHelper::take_snapshot error bad status", 'CRIT')
            raise Exception("VMWareHelper::take_snapshot", status)

    def delete_snapshot(self):
        self._log.log("VMWareHelper:: try to delete snapshot", 'DEBUG')
        self._log.log("Deleting last snapshot")
        onoffStream = os.popen("%s deletesnapshot \"%s\"" \
                                % (self._vmwarerunpath, self._vmxpath))
        status = onoffStream.read().strip()
        if status != '':
            self._log.log("VMWareHelper::delete_snapshot error bad status", 'CRIT')
            raise Exception("VMWareHelper::delete_snapshot", status)
    
    def start(self):
        self._log.log("VMWareHelper:: try to start VM", 'DEBUG')
        if self.get_status() == 'on':
            self._log.log("VMWareHelper:: status already on", 'DEBUG')
            return True
        
        self._log.log('Try to start the vm')
        onoffStream = os.popen("%s \"%s\" start| cut -d \" \" -f 3" \
                               % (self._vmwarecmdpath, self._vmxpath))
        time.sleep(2)
        status = onoffStream.read().strip()
        if status != '1':
            self._log.log("VMWareHelper::start error bad status", 'CRIT')
            raise Exception('VWareHelper::start', 'Command start failed')
        
        self._wait_for_status('on', self._wait_timeout)
        return self.get_status() == 'on'

    def stop(self):
        self._log.log("VMWareHelper:: try to stop VM", 'DEBUG')
        if (self.get_status() == 'off') or (self.get_status() == 'suspended'):
            self._log.log("VMWareHelper:: status already off", 'DEBUG')
            return True
        
        self._log.log('Try to stop the vm')
        onoffStream = os.popen("%s \"%s\" stop trysoft| cut -d \" \" -f 3"\
                                % (self._vmwarecmdpath, self._vmxpath))
        time.sleep(2)
        status = onoffStream.read().strip()
        
        self._wait_for_status('off', self._wait_timeout)
        
        if self.get_status() != 'off':
            self._log.log('Hard stop the vm')
            onoffStream = os.popen("%s \"%s\" stop hard| cut -d \" \" -f 3"\
                                    % (self._vmwarecmdpath, self._vmxpath))
            time.sleep(2)
            status = onoffStream.read().strip()
            if status != '1':
                self._log.log("VMWareHelper::stop error bad status", 'CRIT')
                raise Exception('VWareHelper::stop', 'Command stop failed')
            
            self._wait_for_status('off', self._wait_timeout)
        
        return self.get_status() == 'off'

    def suspend(self):
        self._log.log("VMWareHelper:: try to suspend VM", 'DEBUG')
        if self.get_status() == 'suspended':
            self._log.log("VMWareHelper:: status already suspended", 'DEBUG')
            return True
        
        self._log.log('Try to suspend the vm')        
        onoffStream = os.popen("%s \"%s\" suspend hard|  cut -d \" \" -f 3"\
                                % (self._vmwarecmdpath, self._vmxpath))
        time.sleep(2)
        status = onoffStream.read().strip()
        if status != '1':
            self._log.log("VMWareHelper::suspend error bad status", 'CRIT')
            raise Exception('VWareHelper::suspend', 'Command suspend failed')
        
        self._wait_for_status('suspended', self._wait_timeout)
        return self.get_status() == 'off'

    def _get_root_parent(self, vmdk_file):
        tmp_path = self._tofullpath(vmdk_file)
        hasparent = True
        while hasparent:
            vmx_data = open(tmp_path, 'r').read(10024)
            pattern = '^parentFileNameHint=\"(.*.\.vmdk)\"'
            res = re.compile(pattern, re.IGNORECASE|re.MULTILINE).search(vmx_data)
            if res:
              tmp_path = self._tofullpath(res.group(1))
            else:
                hasparent = False
        return tmp_path
        
    def get_vmdk_list(self):
        vmx_data = open(self._vmxpath, 'r').read(10024)
        pattern = '^.*=.*\"(.*.\.vmdk)\".*'
        rx = re.compile(pattern, re.IGNORECASE|re.MULTILINE)
        vmdk_files = rx.findall(vmx_data)
        return vmdk_files
    
    def _get_filelist_byext(self, file_extlist):                                 
        fileList = [os.path.normcase(f)
                    for f in os.listdir(self._vmrootpath)]           
        fileList = [os.path.join(self._vmrootpath, f) 
                   for f in fileList
                    if os.path.splitext(f)[1] in file_extlist]
        return fileList
    
    def _get_inner_vmdk_list(self, vmdk_file):
        self._log.log("Enter in _get_inner_vmdk_list: %s" % vmdk_file, 'DEBUG')
        inner_list = []
        vmdk_data = open(vmdk_file, 'r').read(10024)
        pattern = '.*\"(.*\.vmdk)\".*'
        rx = re.compile(pattern, re.IGNORECASE|re.MULTILINE)
        files = rx.findall(vmdk_data)
        for file in files:
            inner_list.append(self._tofullpath(file))
        
        return inner_list
    
    def buid_fileslist(self, copy_mode):
        all_files = []
        vmdk_list = []
        tmp_vmdk_list = self.get_vmdk_list()
        self._log.log("VMWareHelper::build_filelist get vmdk report: %s" % tmp_vmdk_list)
        
        if copy_mode == 'HOT':
            for vmdk in tmp_vmdk_list:
                vmdk_list.append(self._get_root_parent(vmdk))
            for vmdk in vmdk_list:
                all_files.extend(self._get_inner_vmdk_list(vmdk))
            all_files.extend(vmdk_list)
            self._log.log("VMWareHelper::build_filelist return with: %s" % all_files, 'DEBUG')
        elif copy_mode == 'COLD':
            all_files.append(self._vmxpath)
            all_files.extend(self._get_filelist_byext(['.vmdk', '.vmsd']))
        elif copy_mode == 'SUSPEND':
            all_files.append(self._vmxpath)
            all_files.append(os.path.join(self._vmrootpath, "nvram"))
            all_files.extend(self._get_filelist_byext(['.vmdk', '.vmsd',
                                                       '.vmem', '.vmss']))            
            
        return all_files

    def _tofullpath(self, relpath):
        buffer = os.path.join(self._vmrootpath, relpath)
        return buffer
    
    def _shell_encode_str(self, buffer):
        return buffer.replace(' ', "\\ ")
    
    def _wait_for_status(self, status, timeout_sec=10):
        status = status.lower()
        start = time.time()
        while (self.get_status().lower() != status) and \
              ((time.time() - start) < timeout_sec):
            self._log.log("still waiting for status to be `status`", 'DEBUG')
            time.sleep(5)
        if (self.get_status().lower() == status):
            return "ok"
        else:
            return "timeout"
        
    def copy(self, file_list, dest_path):
        if not os.path.exists(dest_path):
            self._log.log('VMWareHelper::copy Dest_path daes not exists', 'DEBUG')
            raise Exception('VMWareHelper::copy', 
                            'dest_path dir does not exists!')
        
        for file in file_list:
            srcname = file
            dstname = os.path.join(dest_path, os.path.basename(file))
            
            try:
                self._log.log("Copy from %s to %s" % (srcname, dstname))
                shutil.copy(srcname, dstname)
                self._log.log("done")
                
            except (IOError, os.error), why:
                self._log.log("Can't copy %s to %s: %s" % (srcname, dstname, str(why)), 'CRIT')
    
    def __init__(self, vmxpath, log, vmwarecmdpath = '/usr/bin/vmware-cmd',
                 vmwarerunpath = '/usr/bin/vmrun'):
        self._log = log
        self._vmxpath = vmxpath
        self._vmxshellpath = self._shell_encode_str(vmxpath)
        self._vmwarecmdpath = vmwarecmdpath
        self._vmwarerunpath = vmwarerunpath
        self._vmrootpath = os.path.dirname(vmxpath)

class _OVIMJob:
    _conn = None
    _job_id = 0
    _cursor = None
    _vmxpath = ''
    _job_type = ''
    _job_mode = ''
    _nfs_rootpath = ''
    _alldisks = True
    _remote_pathname = ''
    _log = None
    _vmxfilename = ''
    _storage_remote_path = ''
    _vmspath = ''
    _vmname = ''
    _tmpdirsuffix = ".tmp.1"
    _tmpswapdirsuffix = ".swap"
    _reco_mode = ''
    
    def __init__(self, mysql_conn, job_id, remote_pathname, debug=False):
        self._conn = mysql_conn
        self._log = _OVIMLog(mysql_conn, 'OVIMJob', job_id, debug, debug)
        self._log.log("New job instance (Job ID: %d)" % job_id, 'DEBUG') 
        self._job_id = job_id
        self._cursor = self._conn.cursor()
        self._remote_pathname = remote_pathname
        self._load()
    
    def __del__(self):
        self._cursor.close()

    def _load(self):
        self._log.log("Enter in OVIMJob::_load", 'DEBUG')
        
        query = "SELECT NFSRootPath, vimMain_vm.path as 'VMPath', \
vimMain_vm.configFileName as 'VMXPath', vimMain_jobs.type, \
vimMain_jobs.id, vimMain_jobs.mode, vimMain_localstorage.remotepath, \
vimMain_vmhost.vm_root_path, vimMain_vm.name, vimMain_jobs.reco_mode \
FROM vimMain_jobs, vimMain_vm, vimMain_vmhost, vimMain_localstorage \
WHERE vimMain_jobs.id = '%s' AND \
vimMain_jobs.vm_id = vimMain_vm.id AND \
vimMain_localstorage.id = vimMain_jobs.storage_id AND \
vimMain_vm.vmhost_id = vimMain_vmhost.id" % (self._job_id)

        self._log.log("Will query db with SQL: %s" % query, 'DEBUG')
        self._cursor.execute(query)
        self._log.log("Query done with row: %d" % self._cursor.rowcount, 'DEBUG')
        if self._cursor.rowcount == 0:
            raise Exception('_OVIMJob::_load', 'empty result')
        
        row = self._cursor.fetchone()
        self._vmxpath = os.path.join(row[1], row[2])
        self._vmxfilename = row[2]
        self._storage_remote_path = row[6]
        self._vmspath = row[7]
        self._reco_mode = row[8]
        self._vmname = row[8]
        self._job_type = row[3]
        self._job_mode = row[5]
        self._nfs_rootpath = row[0]
        
        self.update_pid()
        
        self._log.log("Load compelete, data:: vmxpath=%s, job_type=%s, job_mode=%s\
        nfs_rootpath=%s" %
        (self._vmxpath, 
        self._job_type, 
        self._job_mode, 
        self._nfs_rootpath), 'DEBUG'
        )
        
    def _doquery(self, query):
        try:
            self._cursor.execute(query)
        except:
            pass
        
    def _dorecover(self):
        self._log.log("Enter in OVIMJob::_dorecover", 'DEBUG')
        self._log.log("Processing recovery")
        if self._reco_mode == 'REP':
            helper = VMWareHelper(self._vmxpath, self._log)
            remotepath = os.path.join(self._storage_remote_path, self._remote_pathname)
            remotevmxpath = os.path.join(remotepath, self._vmxfilename)
            self._log.log("Try a HOT recovery from %s " % remotevmxpath)
            
            self._log.log("Create local recovery tmp directory")
            localtmpdir = os.path.join(self._vmspath, "%s%s" % (self._vmname, self._tmpdirsuffix))
            if os.path.isdir(localtmpdir):
                self._log.log("Tmp directory %s already exists remove it first!" % localtmpdir, 'DEBUG')
                os.removedirs(localtmpdir)
            
            self._log.log("Create local tmp directory: %s" % localtmpdir) 
            os.makedirs(localtmpdir)
            
            helper_reco = VMWareHelper(remotevmxpath, self._log)
            all_files = helper_reco.buid_fileslist('COLD')
            all_files.append(remotevmxpath)
            
            self._log.log("Buildind File List  Completed: %s" % all_files, 'DEBUG')
            self._log.log("Starting copy...")
            helper_reco.copy(all_files, localtmpdir)
            self._log.log("Stop the VM to do the recovery swap")
            helper.stop()
            self._log.log("Done.")
            origpath = os.path.normpath(os.path.dirname(self._vmxpath))
            swappath = "%s%s" % (origpath, datetime.datetime.now().strftime("-reco-%Y-%m-%d_%H:%M"))
            self._log.log("Renaming %s to %s" % (origpath, swappath))
            os.rename(origpath, swappath)
            self._log.log("Done.")
            self._log.log("Renaming %s to %s" % (localtmpdir, origpath))
            os.rename(localtmpdir, origpath)
            self._log.log("Done.")
            self._log.log("The previous instance of this VM had been renamed and is to be deleted manually (%s)" % swappath)
            self._log.log("Recovery complete")
            
        elif self._reco_mode == 'MVE':
            origpath = os.path.normpath(os.path.dirname(self._vmxpath))
            remotepath = os.path.join(self._storage_remote_path, self._remote_pathname)
            remotevmxpath = os.path.join(remotepath, self._vmxfilename)
            self._log.log("Try a MOVE recovery from %s " % remotevmxpath)
            if os.path.isdir(origpath):
                raise Exception("A VM with that name already exists please remove the complete directory before atempting the recovery!!")
            
            self._log.log("Create the local recovery directory: %s" % origpath)
            os.makedirs(origpath)
            helper_reco = VMWareHelper(remotevmxpath, self._log)
            all_files = helper_reco.buid_fileslist('COLD')
            all_files.append(remotevmxpath)
            self._log.log("Buildind File List  Completed: %s" % all_files, 'DEBUG')
            self._log.log("Starting copy...")
            helper_reco.copy(all_files, origpath)
            self._log.log("Done.")
            
            self._log.log("Recovery complete, you must register the recovered VM by yourself in the directory: %s" % origpath)
            
        else:
            raise Exception("_OVIMJob::recovery unhandled recovery mode : %s" % self._reco_mode)
        
        
    def _dobackup(self):
        self._log.log("Enter in OVIMJob::_dobackup", 'DEBUG')
        self._log.log("Processing backup")
        helper = VMWareHelper(self._vmxpath, self._log)
        
        self._log.log("Check destination")
        if not os.path.isdir(os.path.join(self._storage_remote_path, self._remote_pathname)):
            self._log.log("Create destination")
            os.makedirs(os.path.join(self._storage_remote_path, self._remote_pathname))
        
        vmstatus = helper.get_status()
        
        if self._job_mode == 'HOT':
            self._log.log("Try a HOT backup")
            self._log.log("Take a snapshot")
            helper.take_snapshot()
            self._log.log("done")
        elif self._job_mode == 'COLD':
            if helper.get_status() == 'on':
                self._log.log("Try to stop VM")
                helper.stop()
                self._log.log("done", )
        elif self._job_mode == 'SUSPEND':
                self._log.log("Try to suspend VM")
                helper.suspend()
                self._log.log("done")

        if self._alldisks:
            all_files = helper.buid_fileslist(self._job_mode)
            self._log.log("Buildind File List  Completed: %s" % all_files, 'DEBUG')
            self._log.log("Starting copy...")
            helper.copy(all_files, os.path.join(self._storage_remote_path, self._remote_pathname))
            self._log.log("Done.")                
        
        if self._job_mode == 'HOT':
            self._log.log("Stop to VM to remove snapshot")
            helper.stop()
            helper.delete_snapshot()
            self._log.log("done")
            self._log.log("Copy VMX to dest")
            helper.copy([self._vmxpath], os.path.join(self._storage_remote_path, self._remote_pathname))
            self._log.log("done")      
        
        if vmstatus == 'on':
            self._log.log("Restarting the VM")
            helper.start()
                
        self._log.log("Backup complete")
        self._log.log("Exit from OVIMJob::_dobackup", 'DEBUG')
    
    def update_state(self, newstate):
        query = "UPDATE vimMain_jobs SET state = '%s' WHERE id = '%s'"\
                 % (newstate, self._job_id)
        self._doquery(query)
        self._log.log("Job sate changed to %s" % newstate)
    
    def update_pid(self):
        query = "UPDATE vimMain_jobs SET job_pid = '%s' WHERE id = '%s'"\
                 % (str(os.getpid()), self._job_id)
        self._doquery(query)
    
    def _set_startdate_now(self):
        query = "UPDATE vimMain_jobs SET startdate = NOW() WHERE id = '%s'"\
                 % (self._job_id)
        self._doquery(query)
        
    def _set_enddate_now(self):
        query = "UPDATE vimMain_jobs SET enddate = NOW() WHERE id = '%s'"\
                 % (self._job_id)
        self._doquery(query)
    
    def update_percentage(self, percent):
        query = "UPDATE vimMain_jobs SET percentage = '%s' WHERE id = '%s'" \
                % (percent, self._job-id)
        self._doquery(query)
    
    def process(self):
        self._log.log("%s Process started." % self._job_type)
        self._set_startdate_now()
        if self._job_type == 'BACKUP':
            self._dobackup()
        elif self._job_type == 'RECOVER':
            self._dorecover()
        self._log.log("Process done.")
        self._set_enddate_now()
    
class OVIMClient:
    _dbhost = ''
    _dbport = 3306
    _dbuser = 'OVIM'
    _dbpasswd = 'OVIM'
    _debug = False
    
    _myhostname = 'localhost'
    _mydomain = ''
    _vmwarecmdpath = '/usr/bin/vmware-cmd'
    _log = None
    
    def _usage(self):
        print """       
################################################################################       
#               OVIMClient                   #
# Vincent NICOLAS (vnt.nicolas@gmail.com)"   #   
#
# Purpose: OVIMClient is aim to work with OVIM Server
#                
# Options:  --dbhost: ip or address of the mysql server
#           --dbport: port of the listening mysql server
#           --dbuser: a username able to read/write the OVIM jobs table
#           --dbpasswd: the according username password (could not be empty)
#           -d: Debug output
################################################################################
              """
    def __del__(self):
        if hasattr(self, '_con'):
            self._conn.close()
            
    def __init__(self, DBHost, DBPort, DBUser, DBPasswd):
        locale.setlocale(locale.LC_ALL,'')
        self._dbhost = DBHost
        self._dbport = DBPort
        self._dbuser = DBUser
        self._dbpasswd = DBPasswd
        self._debug = True

    def _parse_args(self):
        try:
            opts, args = getopt.getopt(sys.argv[1:], 'd', \
                                       ['dbhost=', 'dbport=', 'dbuser=', \
                                        'dbpasswd='])
            for o, a in opts:
                if o in ("--dbhost"):
                    self._dbhost = a
                if o in ("--dbport"):
                    self._dbport = a
                if o == "--dbuser":
                    self._dbuser = a
                if o == "--dbpasswd":
                    self._dbpasswd = a
                if o == "-d":
                    self._debug = True

        except getopt.GetoptError:
            self._usage()
            sys.exit(2)

    def _connect_to_db(self):
        try:
            self._conn = MySQLdb.connect(host = self._dbhost, 
                                         user = self._dbuser, 
                                         passwd = self._dbpasswd, 
                                         db = "OVIM")

        except MySQLdb.Error, e:
            print "Error while connecting to MySQL, %d: %s"\
                    % (e.args[0], e.args[1])
            sys.exit (2)
    
    def _get_hostname(self):
        import socket        
        self._myhostname = socket.gethostname().split('.')[0]
    
    def _get_new_job(self):
        self._log.log("Enter in OVIMClient::_get_new_job", 'DEBUG')
        query = "SELECT vimMain_jobs.id, vimMain_vm.name FROM vimMain_jobs, \
vimMain_vm, vimMain_vmhost WHERE (vimMain_vmhost.name = '%s') AND \
(vimMain_jobs.state = 'READY') AND (vimMain_jobs.vm_id = vimMain_vm.id) \
AND (vimMain_vm.vmhost_id = vimMain_vmhost.id) LIMIT 1" % self._myhostname
        
        self._log.log("Will query db with : %s" % query, 'DEBUG')
        cursor = self._conn.cursor()
        cursor.execute (query)
        self._log.log("query done, row count %d" % cursor.rowcount, 'DEBUG')
        if cursor.rowcount > 0:
            row = cursor.fetchone()
            cursor.close()
            self._log.log("Will create a new job instance with args %s, %s/%s" % (row[0], self._myhostname, row[1]), 'DEBUG')
            newjob = _OVIMJob(self._conn, row[0], "%s/%s" % (self._myhostname, row[1]), True)
            self._log.log("Exit from OVIMClient::_get_new_job with a newjob", 'DEBUG')
            return newjob
        else:
            self._log.log("Exit from OVIMClient::_get_new_job with No newjob", 'DEBUG')
            return None

    def _initialize(self):
        self._get_hostname()
        self._connect_to_db()
        self._log = _OVIMLog(self._conn, 'OVIMClient', 0, self._debug, self._debug)                            
        self._log.log("Exit from OVIMClient::initialize", 'DEBUG')
    
    def main(self):
        self._initialize()
        curjob = self._get_new_job()
        while curjob != None:
            self._log.log("New job retreive start processing", 'INFO')
            try: 
                try:
                    curjob.update_state('PERFORMING')
                    curjob.process()
                    curjob.update_state('SUCCESS')
                    
                except Exception, err:
                    self._log.log("ERROR DETECTED IN JOB LOOP: %s" % err, 'CRIT')
                    curjob.update_state('ERROR')
                    curjob._set_enddate_now()
                    self._log.log("Unexpected error: %s" % err)
            
            finally:
                del curjob
                self._log.log("Job process done.", 'INFO')
                curjob = self._get_new_job()
                
        self._log.log("OVIMClient shutdown", 'INFO')
        
if __name__ == "__main__":
    client = OVIMClient('srvnas', '3306', 'OVIM', 'OVIM123')
    client.main()